from dwave.system.composites import EmbeddingComposite
# from dwave.system.samplers import DWaveSampler
import neal
import dwavebinarycsp
from penalty_function import makeExactOne, gridSweep
from collections import defaultdict
import dimod
import random

# 节点
# 时间

TIME_FRAME_NUMBER = 10  # 假设 0.1 us 一段
TIME_FRAMES = [i for i in range(TIME_FRAME_NUMBER)]

NODE_NUMBER = 5
NODES = [i for i in range(NODE_NUMBER)]


QUEUE_SIZE_PER_NODE = {
    node: int(random.random() * 5)
    for node in NODES    
}

AVG_EXECUTION_TIME_PER_BODE = {
    node: random.random() + 0.0001
    for node in NODES    
}

# WAIT_TIME_PAR_NODE = {
#     node: node % 2
#     for node in NODES    
# }  #queue size * average execute time


# 初始化网速
MAX_SPEED = 10
NETWORK_SPEED = {
    (node1, node2): (node1 + node2) % 3 + 1
    for node1 in NODES
    for node2 in NODES    
}
for node in NODES:
    NETWORK_SPEED[(node, node)] = 10

# node读取时间、remote node节点的时间
def getNodeQubit(own_node, read_node, time_frame):
    return f'Node{own_node}_Read{read_node}_Time{time_frame}'

sampler = neal.SimulatedAnnealingSampler()

# 每个节点的queue的数量 * 平均时间 要小
# 每个节点的利用率高（1的数量多）

# 每次都的时候要考虑 （传输速度，等待时间）,等待时常是考虑了之前的read的
# 每个节点应该发的节点是exact one的（未来可以试下有多个node，自己算一个传输快的）
Q = defaultdict(int)

# # 每次只能有一个节点
exact_one_lagrange = 5  #要是太小了就没办法变成硬constraint了，以后要自适应下
for time in TIME_FRAMES:
    for own_node in NODES:  #实际持有qubit的node
        qubits = [getNodeQubit(own_node, read_node, time,) for read_node in NODES]
        makeExactOne(Q, qubits, exact_one_lagrange)

# # 必须得在最快被响应
# # 假设数据大小固定 (queue size * average execution time + data size / network size)

time_cost_lagrange = 0.1
DATA_SIZE = 5
for time in TIME_FRAMES:
    for own_node in NODES:  #实际持有qubit的node
        for read_node in NODES:
            translate_speed = NETWORK_SPEED[(own_node, read_node)]

            queue_size = QUEUE_SIZE_PER_NODE[read_node]
            avg_execution_time = AVG_EXECUTION_TIME_PER_BODE[read_node]
            response_time = queue_size * avg_execution_time +  DATA_SIZE / translate_speed

            qubit = getNodeQubit(own_node, read_node, time)
            Q[(qubit, qubit)] += response_time * time_cost_lagrange

# 互相之间尽量不要发到一起，如果一样，总时长会加avg_execution_time
for time in TIME_FRAMES:
    for read_node in NODES:
        for own_node1, own_node2 in gridSweep(NODES):
            qubit1 = getNodeQubit(own_node1, read_node, time)
            qubit2 = getNodeQubit(own_node2, read_node, time)
            avg_execution_time = AVG_EXECUTION_TIME_PER_BODE[read_node]
            Q[(qubit1, qubit2)] += avg_execution_time * time_cost_lagrange

# # 如果一个node被add_queue了（也就是为1），后面的对应avg_execution_time的node被选中的概率都会被depress
for time in TIME_FRAMES:
    for own_node in NODES:
        for read_node in NODES:
            now_read_qubit = getNodeQubit(own_node, read_node, time)

            avg_execution_time = AVG_EXECUTION_TIME_PER_BODE[read_node]
            effect_times = TIME_FRAMES[time+1: time+int(avg_execution_time)]  #被前面影响的时间
            
            for effect_own_node in NODES:
                for effect_time in effect_times:
                    effect_read_qubit = getNodeQubit(effect_own_node, read_node, effect_time)
                    Q[(now_read_qubit, effect_read_qubit)] += avg_execution_time * time_cost_lagrange   # now_read_qubit * avg_execution_time * later_read_qubit  如果前一个已经被分配了一个queue task那后一个的等待时间就会加一个avg_execution_time

# 需不需要加入一个goal，今早算完：不需要，因为不知道有多少任务（越早开始是否表示越早结束）

# 假设当前定了，状态会变成这个定的


# csp = dwavebinarycsp.ConstraintSatisfactionProblem(dwavebinarycsp.BINARY)

# bqm = dwavebinarycsp.stitch(csp)

# https://link.springer.com/article/10.1007/s10951-019-00629-3
# online schedule algorithm

bqm = dimod.BQM.from_qubo(Q) # (A, B) and (B, A) will be mergeed automatically

sampler = neal.SimulatedAnnealingSampler()
results = sampler.sample(bqm, num_reads=100)

def praiseQubitName(qubit_name):
    own_node, read_node, time = qubit_name.split('_')

    return int(own_node.replace('Node', '')), int(read_node.replace('Read', '')), int(time.replace('Time', ''))

def praise2frame_qubit(sample):
    frame2node2assign_node = [
        [
            None
            for _ in NODES
        ]
        for _ in TIME_FRAMES
    ]
    # 麻了，乘号里面竟然是引用复制

    for qubit_name, value in sample.items():
        own_node, read_node, time = praiseQubitName(qubit_name)

        # frame2node2read_node[time][own_node][read_node] = value

        if value == 1:
            assert frame2node2assign_node[time][own_node] == None
            frame2node2assign_node[time][own_node] = read_node
    

    return frame2node2assign_node

for sample, energy, num_occurrences in list(results.data(['sample', 'energy', 'num_occurrences'])):
    # print(sample, energy, num_occurrences)
    # print(energy)

    frame2node2assign_node = praise2frame_qubit(sample)

    pass